package A一周刷爆LeetCode.A基础班.H前缀树和贪心算法;

import org.apache.commons.lang3.StringUtils;

/**
 * @author wei.zhao
 * @description: 前缀树
 * @date 2022年03月24日 15:58
 */
public class Code01_TrieTree {

    /* 前缀树节点 */
    public static class TrieNode {
        public int pass;
        public int end;
        public TrieNode[] nexts;

        public TrieNode() {
            nexts = new TrieNode[26];
        }
    }

    /* 前缀树 */
    public static class Trie {
        private TrieNode root;

        public Trie() {
            root = new TrieNode();
        }

        /* 插入字符串 */
        public void insert(String word) {
            if (StringUtils.isBlank(word)) {
                return;
            }
            root.pass++;
            char[] chars = word.toCharArray();
            TrieNode cur = root;
            for (char c : chars) {
                if (cur.nexts[c - 'a'] == null) {
                    // 注意：不要忘记将新增的节点赋值到数组中
                    // 小技巧：先确定数组中的值，后面再统一赋值
                    cur.nexts[c - 'a'] = new TrieNode();
                }
                cur = cur.nexts[c - 'a'];
                cur.pass++;
            }
            cur.end++;
        }

        /* 删除字符串 */
        public void delete(String word) {
            if (search(word) == 0) {
                return;
            }
            root.pass--;
            TrieNode cur = root;
            char[] chars = word.toCharArray();
            for (char c : chars) {
                TrieNode next = cur.nexts[c - 'a'];
                if (--next.pass == 0) {
                    cur.nexts[c - 'a'] = null;
                }
                cur = next;
            }
            cur.end--;
        }

        /* 查找某个字符串个数 */
        public int search(String word) {
            if (StringUtils.isBlank(word)) {
                return 0;
            }
            char[] chars = word.toCharArray();
            TrieNode cur = root;
            for (char c : chars) {
                cur = cur.nexts[c - 'a'];
                if (cur == null) {
                    return 0;
                }
            }
            return cur.end;
        }

        /* 查找以某个字符串开头的字符串数 */
        public int prefixNumber(String pre) {
            if (StringUtils.isBlank(pre)) {
                return 0;
            }
            char[] chars = pre.toCharArray();
            TrieNode cur = root;
            for (char c : chars) {
                cur = cur.nexts[c - 'a'];
                if (cur == null) {
                    return 0;
                }
            }
            return cur.pass;
        }
    }

    public static void main(String[] args) {
        Trie trie = new Trie();
        System.out.println(trie.search("zuo"));
        trie.insert("zuo");
        System.out.println(trie.search("zuo"));
        trie.delete("zuo");
        System.out.println(trie.search("zuo"));
        trie.insert("zuo");
        trie.insert("zuo");
        trie.delete("zuo");
        System.out.println(trie.search("zuo"));
        trie.delete("zuo");
        System.out.println(trie.search("zuo"));
        trie.insert("zuoa");
        trie.insert("zuoac");
        trie.insert("zuoab");
        trie.insert("zuoad");
        trie.delete("zuoa");
        System.out.println(trie.search("zuoa"));
        System.out.println(trie.prefixNumber("zuo"));
    }

}
